import os
import tensorflow as tf
from tensorflow.python.keras import Sequential, Model, layers, optimizers, losses
import keras_preprocessing.sequence
from keras.datasets import imdb

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

batches = 128  # 批量大小
total_words = 10000  # 词汇表大小
max_review_len = 80  # 句子的最大长度,大于的部分将截断，小于的将填充
embedding_len = 100  # 词向量特征长度

# 加载IMDB数据集，数据采用数字编码，一个数字代表一个单词
(x_train, y_train), (x_test, y_test) = imdb.load_data(num_words=total_words)  # 单词数量限制，假设只有常用一万个单词

# 截断和填充句子，使得等长为max_review_len，长句子保留后面部分，短句子在前面填充
x_train = keras_preprocessing.sequence.pad_sequences(x_train, maxlen=max_review_len)
x_test = keras_preprocessing.sequence.pad_sequences(x_test, maxlen=max_review_len)

# 构建数据集，打散，批量，并丢掉最后一个不够batches的batch
db_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))
db_train = db_train.shuffle(1000).batch(batches, drop_remainder=True)

db_test = tf.data.Dataset.from_tensor_slices((x_test, y_test))
db_test = db_test.batch(batches, drop_remainder=True)
# 25k个句子，每个句子80个单词，   y---1 好评   y---0差评
print('x_train shape:', x_train.shape, tf.reduce_max(y_train), tf.reduce_min(y_train))
print('x_test shape:', x_test.shape)


class MyRNN(Model):
    def __init__(self, units):
        super(MyRNN, self).__init__()
        # 词向量编码 [b,80] ==> [b,80,100]   每个单词用长度100的向量来表示
        # embedding_len：单词向量的长度，total_words：词汇的数量 max_review_len：输入句子长度
        self.embedding = layers.Embedding(total_words, embedding_len, input_length=max_review_len)
        # 构建2个Cell [b,80,100] => [b,64]
        self.rnn = Sequential([
            layers.LSTM(units, dropout=0.5, return_sequences=True),
            layers.LSTM(units, dropout=0.5)

            # layers.SimpleRNN(units, dropout=0.5, return_sequences=True),
            # layers.SimpleRNN(units, dropout=0.5)
        ])
        # 构建分类网络，用于将Cell的输出特征进行分类，2分类
        # [b,64]=> [b,1]
        self.outlayer = Sequential([
            layers.Dense(units),
            layers.Dropout(0.5),
            layers.ReLU(),
            layers.Dense(1)
        ])

    # 向前计算
    def call(self, inputs, training=None):
        x = inputs  # [b,80]
        # embedding: [b,80] ==> [b,80, 100]
        x = self.embedding(x)
        # rnn-全连接神经网络 cell compute: [b, 80,100] => [b,64]
        out1 = self.rnn(x)
        # 末层最后一个输出作为分类网络的输入: [b, 64] => [b, 1]
        x = self.outlayer(out1, training)
        # p(y is pos|x)
        prob = tf.sigmoid(x)
        return prob


def main():
    units = 64  # rnn状态向量长度
    epochs = 5
    model = MyRNN(units)
    # 装配优化器，学习率，测量器
    model.compile(optimizer=optimizers.adam_v2.Adam(1e-3),
                  loss=losses.BinaryCrossentropy(),
                  metrics=['accuracy'])

    # 训练 与 验证  validation_data：验证数据
    model.fit(db_train, epochs=epochs, validation_data=db_test)
    # 测试
    model.evaluate(db_test)


if __name__ == '__main__':
    main()
